from sentence_transformers import CrossEncoder

from app.decorator.timeit import timeit


# rerank 模型,通过transform
@timeit
def init_rerank_model():
    return CrossEncoder(model_name_or_path="./transformers/bge-reranker-large", local_files_only=True,
                        cache_folder="./transformers/.cache")


bge_rerank_model = init_rerank_model()

# 例子
# query = "Java代码如何进行性能优化"
# 
# document = [
#     "SpringCloud微服务从入门到精通",
#     "元宇宙与数字经济",
#     "Java性能优化"
# ]
# result = rerank_model.rank(query=query, documents=document, return_documents=True, top_k=3, show_progress_bar=True)
# print(result)
